{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "d = 4                                           # 向量维度\n",
    "nb = 10                                      # index向量库的数据量\n",
    "nq = 2                                       # 待检索query的数目\n",
    "np.random.seed(1234)             \n",
    "xb = np.random.random((nb, d)).astype('float32')\n",
    "xb[:, 0] += np.arange(nb) / 1000.                # index向量库的向量\n",
    "xq = np.random.random((nq, d)).astype('float32')\n",
    "xq[:, 0] += np.arange(nq) / 1000.                # 待检索的query向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "base:\n",
      " [[0.19151945 0.62210876 0.43772775 0.7853586 ]\n",
      " [0.7809758  0.2725926  0.27646425 0.8018722 ]\n",
      " [0.96013933 0.87593263 0.35781726 0.5009951 ]\n",
      " [0.68646294 0.71270204 0.37025076 0.5611962 ]\n",
      " [0.5070832  0.01376845 0.7728266  0.8826412 ]\n",
      " [0.36988598 0.6153962  0.07538124 0.368824  ]\n",
      " [0.9391401  0.65137815 0.39720258 0.78873014]\n",
      " [0.32383612 0.56809866 0.8691274  0.4361734 ]\n",
      " [0.81014764 0.14376682 0.70426095 0.7045813 ]\n",
      " [0.22779211 0.92486763 0.44214076 0.90931594]]\n",
      "query:\n",
      " [[0.05980922 0.18428709 0.04735528 0.6748809 ]\n",
      " [0.59562474 0.5333102  0.04332406 0.5614331 ]]\n"
     ]
    }
   ],
   "source": [
    "print(\"base:\\n\", xb[:10])\n",
    "print(\"query:\\n\", xq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8.0\n",
      "False\n",
      "True\n",
      "Train time: 0.059209346771240234\n",
      "Add time: 0.0024623870849609375\n",
      "None\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING clustering 10 points to 3 centroids: please provide at least 117 training points\n"
     ]
    }
   ],
   "source": [
    "import faiss\n",
    "print(faiss.__version__)    \n",
    "dim, measure = d, faiss.METRIC_L2\n",
    "param = 'IVF3,Flat'\n",
    "index = faiss.index_factory(dim, param, measure)\n",
    "print(index.is_trained)\n",
    "st = time.time()\n",
    "index.train(xb)\n",
    "et = time.time()\n",
    "print(index.is_trained)\n",
    "print(\"Train time:\",et-st)\n",
    "index.add(xb)\n",
    "eet = time.time()\n",
    "print(\"Add time:\",eet-et)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "10\n"
     ]
    }
   ],
   "source": [
    "print(index.is_trained)         # 输出为True，代表该类index不需要训练，只需要add向量进去即可\n",
    "print(index.ntotal)             # 输出index中包含的向量总数，为100000 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_bytes_as_hex(data, bytes_per_line=32):\n",
    "    count = 0\n",
    "    for byte in data:\n",
    "        print(hex(byte)[2:].zfill(2), end=' ')\n",
    "        count += 1\n",
    "        if(count==bytes_per_line):\n",
    "            print()\n",
    "            count = 0\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xb\n",
      "[[0.19151945 0.62210876 0.43772775 0.7853586 ]\n",
      " [0.7809758  0.2725926  0.27646425 0.8018722 ]\n",
      " [0.96013933 0.87593263 0.35781726 0.5009951 ]\n",
      " [0.68646294 0.71270204 0.37025076 0.5611962 ]\n",
      " [0.5070832  0.01376845 0.7728266  0.8826412 ]\n",
      " [0.36988598 0.6153962  0.07538124 0.368824  ]\n",
      " [0.9391401  0.65137815 0.39720258 0.78873014]\n",
      " [0.32383612 0.56809866 0.8691274  0.4361734 ]\n",
      " [0.81014764 0.14376682 0.70426095 0.7045813 ]\n",
      " [0.22779211 0.92486763 0.44214076 0.90931594]]\n",
      "ad 1d 44 3e 85 42 1f 3f da 1d e0 3e 43 0d 49 3f \n",
      "08 ee 47 3f 42 91 8b 3e b9 8c 8d 3e 7f 47 4d 3f \n",
      "b1 cb 75 3f 1f 3d 60 3f d3 33 b7 3e 37 41 00 3f \n",
      "09 bc 2f 3f a4 73 36 3f 82 91 bd 3e 8e aa 0f 3f \n",
      "34 d0 01 3f 10 95 61 3c f7 d7 45 3f c6 f4 61 3f \n",
      "b2 61 bd 3e 9b 8a 1d 3f 7b 61 9a 3d 80 d6 bc 3e \n",
      "7c 6b 70 3f b8 c0 26 3f 23 5e cb 3e 38 ea 49 3f \n",
      "d9 cd a5 3e ea 6e 11 3f 22 7f 5e 3f 1f 52 df 3e \n",
      "d6 65 4f 3f 9c 37 13 3e 72 4a 34 3f 71 5f 34 3f \n",
      "56 42 69 3e 20 c4 6c 3f 46 60 e2 3e ee c8 68 3f \n",
      "\n",
      "xq\n",
      "[[0.05980922 0.18428709 0.04735528 0.6748809 ]\n",
      " [0.59562474 0.5333102  0.04332406 0.5614331 ]]\n",
      "84 fa 74 3d c1 b5 3c 3e 9c f7 41 3d ff c4 2c 3f \n",
      "dd 7a 18 3f 04 87 08 3f 93 74 31 3d 14 ba 0f 3f \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"xb\")\n",
    "print(xb)\n",
    "print_bytes_as_hex(xb.tobytes(),16)\n",
    "print(\"xq\")\n",
    "print(xq)\n",
    "print_bytes_as_hex(xq.tobytes(),16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "D:\n",
      " [[3.7363139e-01 3.7645897e-01 7.8749317e-01 9.4931221e-01]\n",
      " [1.4731415e-01 3.0883682e-01 3.5281977e-01 3.4028235e+38]]\n",
      "I:\n",
      " [[ 0  5  9  7]\n",
      " [ 3  6  2 -1]]\n",
      "[0.19151945 0.62210876 0.43772775 0.7853586 ]\n",
      "ad 1d 44 3e 85 42 1f 3f da 1d e0 3e 43 0d 49 3f \n",
      "Search time: 0.0016078948974609375\n"
     ]
    }
   ],
   "source": [
    "k = 4                     # topK的K值\n",
    "st = time.time()\n",
    "D, I = index.search(xq, k)# xq为待检索向量，返回的I为每个待检索query最相似TopK的索引list，D为其对应的距离\n",
    "et = time.time()\n",
    "print(\"D:\\n\",D)\n",
    "print(\"I:\\n\",I)\n",
    "print(xb[I[0][0]])\n",
    "print_bytes_as_hex(xb[I[0][0]].tobytes())\n",
    "\n",
    "print(\"Search time:\",et-st)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "#持久化索引\n",
    "faiss.write_index(index, \"test1.index\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "#读取索引\n",
    "read_index = faiss.read_index(\"test1.index\")  # 若已知维度，可传入具体值"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "IndexIVFFlat\n"
     ]
    }
   ],
   "source": [
    "#打印索引的信息\n",
    "print(read_index.d)\n",
    "print(type(read_index).__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Search time: 0.0003647804260253906\n",
      "[[14 -1 -1 -1]\n",
      " [ 8 13  4 16]]\n",
      "[[3.0878145e-01 3.4028235e+38 3.4028235e+38 3.4028235e+38]\n",
      " [2.8407644e-02 5.9989795e-02 6.1861880e-02 2.1397224e-01]]\n"
     ]
    }
   ],
   "source": [
    "k = 4                     # topK的K值\n",
    "st = time.time()\n",
    "D, I = read_index.search(xq, k)# xq为待检索向量，返回的I为每个待检索query最相似TopK的索引list，D为其对应的距离\n",
    "et = time.time()\n",
    "print(\"Search time:\",et-st)\n",
    "print(I[:5])\n",
    "print(D[-5:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
